from solvers.solverbase import SolverBase
from solvers.dg.dgbase import DGBase
from dolfin import *
import numpy

class GaussLegendre4(SolverBase,DGBase):
    '2-symplectic method rewritten without increments. DG.'
    def __init__(self,p=1,inflowBC=0):
        SolverBase.__init__(self,p)
        DGBase.__init__(self,inflowBC)

    def solve(self,problem):
        mesh = problem.domain

        def get_pk_m(phiC,phi_,out):
            out.vector()[:] = 0.5*phi_.vector().array() + (0.5-sqrt(3.)/6.)*phiC.vector().array()
            
        def get_pk_p(phiC,phi_,out):
            out.vector()[:] = (0.5-2*sqrt(3.)/6.)*phi_.vector().array() + (0.5+sqrt(3.)/6.)*phiC.vector().array()
        
        def get_sk_m(phiP,phi_,out):
            out.vector()[:] = 0.5*phi_.vector().array() + sqrt(3.)/6.*phiP.vector().array()

        def get_sk_p(phiP,phi_,out):
            out.vector()[:] = (0.5-2*sqrt(3.)/6.)*phi_.vector().array() + sqrt(3.)/6.*phiP.vector().array() 
        
        p_degree = self.p 
        V = FunctionSpace(mesh,'DG',p_degree)
        VV = VectorFunctionSpace(mesh,'DG',p_degree)

        phi_ = interpolate(problem.phi_,V)
        phiP = Function(V)                  # solution to the auxiliary equation
        phiC = Function(V)                  # solution to the true equation
        phi = TrialFunction(V)          
        varphi = TestFunction(V)

        t = problem.t          
        dt = Constant(SolverBase._get_time_step(self,problem,VV,t))
        T = problem.T

        v = SolverBase._get_velocity(self,problem,VV,t)
        pk_m = Function(V)
        pk_p = Function(V)
        sk_m = Function(V)
        sk_p = Function(V)

        n = FacetNormal(mesh)

        if not self.inflowBC:
            aP = phi*varphi*dx 

            lP = phi_*varphi*dx\
                + 0.5*(dt*inner(v,nabla_grad(varphi))*pk_m*dx\
                       - dt('+')*dot(v('+'),jump(varphi,n)*avg(pk_m))*dS\
                       - dt('+')*0.5*abs(dot(v('+'),n('+')))*dot(jump(pk_m,n),jump(varphi,n))*dS\
                       - dt*pk_m*varphi*dot(v,n)*ds\
                      )\
                - 0.5*(dt*inner(v,nabla_grad(varphi))*pk_p*dx\
                       - dt('+')*dot(v('+'),jump(varphi,n)*avg(pk_p))*dS\
                       - dt('+')*0.5*abs(dot(v('+'),n('+')))*dot(jump(pk_p,n),jump(varphi,n))*dS\
                       - dt*pk_p*varphi*dot(v,n)*ds\
                      )

            #-----------------------

            aC = phi*varphi*dx\
                - 0.5*(dt*inner(v,nabla_grad(varphi))*(0.5-sqrt(3.)/6.)*phi*dx\
                       - dt('+')*dot(v('+'),jump(varphi,n)*avg((0.5-sqrt(3.)/6.)*phi))*dS\
                       - dt('+')*0.5*abs(dot(v('+'),n('+')))*dot(jump((0.5-sqrt(3.)/6.)*phi,n),jump(varphi,n))*dS\
                       - dt*(0.5-sqrt(3.)/6.)*phi*varphi*dot(v,n)*ds\
                      )\
                - 0.5*(dt*inner(v,nabla_grad(varphi))*(0.5+sqrt(3.)/6.)*phi*dx\
                       - dt('+')*dot(v('+'),jump(varphi,n)*avg((0.5+sqrt(3.)/6.)*phi))*dS\
                       - dt('+')*0.5*abs(dot(v('+'),n('+')))*dot(jump((0.5+sqrt(3.)/6.)*phi,n),jump(varphi,n))*dS\
                       - dt*(0.5+sqrt(3.)/6.)*phi*varphi*dot(v,n)*ds\
                      )

            lC = phi_*varphi*dx\
                + 0.5*(dt*inner(v,nabla_grad(varphi))*sk_m*dx\
                       - dt('+')*dot(v('+'),jump(varphi,n)*avg(sk_m))*dS\
                       - dt('+')*0.5*abs(dot(v('+'),n('+')))*dot(jump(sk_m,n),jump(varphi,n))*dS\
                       - dt*sk_m*varphi*dot(v,n)*ds\
                      )\
                + 0.5*(dt*inner(v,nabla_grad(varphi))*sk_p*dx\
                       - dt('+')*dot(v('+'),jump(varphi,n)*avg(sk_p))*dS\
                       - dt('+')*0.5*abs(dot(v('+'),n('+')))*dot(jump(sk_p,n),jump(varphi,n))*dS\
                       - dt*sk_p*varphi*dot(v,n)*ds\
                      )
        else:
            (ib,ib_value,bParts) = problem.get_inflowBCS()  # ib_value has bc conditions on the inflow part
            g = Function(V)                                 # inflow bcs for the weak form
 
            aP = phi*varphi*dx 

            lP = phi_*varphi*dx\
                + 0.5*(
                       - dt*inner(v,nabla_grad(pk_m))*varphi*dx\
                       + dt('+')*dot(v('+'),jump(pk_m,n)*avg(varphi))*dS\
                       - dt('+')*0.5*abs(dot(v('+'),n('+')))*dot(jump(pk_m,n),jump(varphi,n))*dS\
                       - dt*(pk_m-g)*varphi*0.5*(abs(dot(v,n))-dot(v,n))*ds\
                      )\
                - 0.5*(
                       - dt*inner(v,nabla_grad(pk_p))*varphi*dx\
                       + dt('+')*dot(v('+'),jump(pk_p,n)*avg(varphi))*dS\
                       - dt('+')*0.5*abs(dot(v('+'),n('+')))*dot(jump(pk_p,n),jump(varphi,n))*dS\
                       - dt*(pk_p-g)*varphi*0.5*(abs(dot(v,n))-dot(v,n))*ds\
                      )

            #-----------------------

            aC = phi*varphi*dx\
                + 0.5*(
                       + dt*inner(v,nabla_grad(phi))*(0.5-sqrt(3.)/6.)*varphi*dx\
                       - dt('+')*dot(v('+'),jump(phi,n)*avg((0.5-sqrt(3.)/6.)*varphi))*dS\
                       + dt('+')*0.5*abs(dot(v('+'),n('+')))*dot(jump((0.5-sqrt(3.)/6.)*phi,n),jump(varphi,n))*dS\
                       + dt*(0.5-sqrt(3.)/6.)*phi*varphi*0.5*(abs(dot(v,n))-dot(v,n))*ds\
                      )\
                + 0.5*(
                       + dt*inner(v,nabla_grad(phi))*(0.5+sqrt(3.)/6.)*varphi*dx\
                       - dt('+')*dot(v('+'),jump(phi,n)*avg((0.5+sqrt(3.)/6.)*varphi))*dS\
                       + dt('+')*0.5*abs(dot(v('+'),n('+')))*dot(jump((0.5+sqrt(3.)/6.)*phi,n),jump(varphi,n))*dS\
                       + dt*(0.5+sqrt(3.)/6.)*phi*varphi*0.5*(abs(dot(v,n))-dot(v,n))*ds\
                      )

            lC = phi_*varphi*dx\
                + 0.5*(
                       - dt*inner(v,nabla_grad(sk_m))*varphi*dx\
                       + dt('+')*dot(v('+'),jump(sk_m,n)*avg(varphi))*dS\
                       - dt('+')*0.5*abs(dot(v('+'),n('+')))*dot(jump(sk_m,n),jump(varphi,n))*dS\
            #           - dt*(sk_m-g)*varphi*0.5*(abs(dot(v,n))-dot(v,n))*ds\
                      )\
                + 0.5*(
                       - dt*inner(v,nabla_grad(sk_p))*varphi*dx\
                       + dt('+')*dot(v('+'),jump(sk_p,n)*avg(varphi))*dS\
                       - dt('+')*0.5*abs(dot(v('+'),n('+')))*dot(jump(sk_p,n),jump(varphi,n))*dS\
             #          - dt*(sk_p-g)*varphi*0.5*(abs(dot(v,n))-dot(v,n))*ds\
                      )\
                +0.5*dt*(0.5-sqrt(3.)/6.)*g*varphi*0.5*(abs(dot(v,n))-dot(v,n))*ds\
                +0.5*dt*(0.5+sqrt(3.)/6.)*g*varphi*0.5*(abs(dot(v,n))-dot(v,n))*ds
        
        self.update(problem,phi_,t,float(dt),0)  # store the solution at time 0

        phiC.assign(phi_)  # initial guess
        
        phiC_ = Function(V) # so that error could be calculated
        phiC_.assign(phi_)

        AP = assemble(aP)
        AC = assemble(aC)
        phi = Function(V)
   
        t = 0.
        count = 0
       
        maxIter = 10
        tol = 1E-8
        
        while t <= T:
            t += float(dt)
            count += 1
            error = 1.

            if self.inflowBC:
                g.vector()[:] = interpolate(ib_value(t),V).vector().array() # update inflow values
      
            iteration = 0

            while error > tol and iteration < maxIter:
                iteration += 1
                get_pk_m(phiC,phi_,pk_m)      #!
                get_pk_p(phiC,phi_,pk_p)      #!
                LP = assemble(lP)
                solve(AP,phi.vector(),LP)  # solve aux system
                phiP.assign(phi)

                get_sk_m(phiP,phi_,sk_m)      #!
                get_sk_p(phiP,phi_,sk_p)      #!
                LC = assemble(lC)
                solve(AC,phi.vector(),LC)  # solve true system
                phiC.assign(phi)

                diff = phiC.vector().array() - phiC_.vector().array()
                error = numpy.linalg.norm(diff,ord=2)        # compute the error
                phiC_.assign(phiC)
                
        #        print iteration, error

            phi_.assign(phiC)
            self.update(problem,phi_,t,float(dt),count)

            
        problem.save_data(self.saveDir,str(self.CFL))
